from common.DataProcessingStep import *
from common.ModelZoo import get_vgg16_model, get_vgg16_model1, get_triplet_model, get_cnn_pay_attention_model
from common.loss_zoo import Contrastive_loss, triplet_semihard_loss,Contrastive_PCC_loss
from common.DataProcessor import JRMMDataProcessor,DataProcessor


def main():
    processor = DataProcessor()
    processor.add_step(Loadxxx())
    processor.add_step(CalculateHammingWeightStep(bit_length=32))
    processor.add_step(PccStep())
    processor.add_step(ExtendedDimensionStep())
    processor.add_step(TrainDLStep(
                                   get_model=get_vgg16_model,
                                   loss=Contrastive_PCC_loss,
                                   epochs=2,
                                   batch_size=256,
                                   learning_rate=5e-4,
                                   output_shape=32,
                                   embedding=True))
    processor.add_step(TrainTAStep())
    processor.add_step(TestTAStep())
    # processor.add_step(TestDLStep())
    processor.process()



if __name__ == '__main__':
    main()
